#!/usr/bin/env python
#########################################
# Author: FuHuizhong<fuhuizn@gmail.com> #
# Licence: GPL3                         #
# Port Mapping Tool with UPNP           #
#########################################
import socket
import thread
import re
import urllib2
import xml.dom.minidom
import time,sys
import base64

class UPNPCoherence:
    def __init__(self):
        self.location = ''
        self.ctrl_url = ''
        self.serviceType = ''
        self.s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.s.settimeout(5)
        self.s.bind( ('',0) )
        #self.s.connect( ('239.255.255.250',1900) )
    #auth_enc = 'YWRtaW46YWRtaW4='
    def newRequest( self,url ):
        req = urllib2.Request( url,unverifiable=True )
        info = urllib2.urlparse.urlparse( url )
        req.add_header( 'HOST', info[1] )
        #global auth_enc
        #req.add_header( 'Authorization', 'Basic ' + auth_enc )
        return req
    
    def is_ok(self,data):
        try:
            s1,s2 = data.split('\r\n',1)
        except:
            s1 = ''
        if s1.endswith('200 OK'):
            return True
        else:
            return False

    def get_location(self,data):
        #print data
        ret = re.findall( 'LOCATION: ([^\r]*)',data )
        location = ret[0]
#        for item in ret:
#            print 'discription:',item
        return location

    def get_ctrlurl(self):
        location = self.location
        fp = urllib2.urlopen( self.newRequest( location ) )
        #print fp.headers
        data = ''
        for buf in fp:
            data += buf
        fp.close()
        #print data
        igd = xml.dom.minidom.parseString(data)
        #get control url from the xml
        for service in igd.getElementsByTagName('service'):
            val = {}
            for child in service.childNodes:
                try:
                    val[child.nodeName.lower()] = child.childNodes[0].nodeValue
                except:
                    pass
            if (val['servicetype'].lower()=='urn:schemas-upnp-org:service:wanpppconnection:1') or (val['servicetype'].lower()=='urn:schemas-upnp-org:service:wanipconnection:1'):
                ctrl_url = urllib2.urlparse.urljoin( location,val['controlurl'] )
                serviceType = val['servicetype']
                break
        #ret = get_port_map( 1 )
        print 'controlurl:',ctrl_url
        #print 'servicetype:',serviceType
        return (ctrl_url,serviceType)

    def get_local_ip(self):
        ds = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        ds.connect(('255.255.255.254', 0))
        self.ip,port = ds.getsockname()
        ds.close()
        return self.ip
 
    def add_port(self,ExternalPort,Protocol,InternalPort):
        ctrl_url = self.ctrl_url
        serviceType = self.serviceType
        if ctrl_url == '':
            return ''
        InternalClient = self.get_local_ip()
        actionName = 'AddPortMapping'
        actionParams = '<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort>%d</NewExternalPort>\r\n\
<NewProtocol>%s</NewProtocol>\r\n\
<NewInternalPort>%d</NewInternalPort>\r\n\
<NewInternalClient>%s</NewInternalClient>\r\n\
<NewEnabled>1</NewEnabled>\r\n\
<NewPortMappingDescription>UPNPCoherence</NewPortMappingDescription>\r\n\
<NewLeaseDuration>0</NewLeaseDuration>\r\n' % (ExternalPort,Protocol,InternalPort,InternalClient)
        control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
        up = urllib2.urlparse.urlparse(ctrl_url)
        header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
        #global auth_enc
        #header += 'Authorization: Basic ' + auth_enc + '\r\n'
        header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
        header += 'CONTENT-TYPE: text/xml\r\n'
        header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
        req = header + control_doc
        host,port = up[1].split(':')
        fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        fd.connect( (host,int(port)) )
        fd.send( req )
        data = ''
        while 1:
            buf = fd.recv(1492)
            if len(buf)>0:
                data += buf
            else:
                break
        fd.close()
#        if self.is_ok(data):
#            print 'OK ADD',InternalClient,Protocol,str(ExternalPort),str(InternalPort)
#        else:
#            print data
        return data

    def del_port(self,ExternalPort,Protocol):
        ctrl_url = self.ctrl_url
        serviceType = self.serviceType
        if ctrl_url == '':
            return ''
        actionName = 'DeletePortMapping'
        actionParams = '<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort>%d</NewExternalPort>\r\n\
<NewProtocol>%s</NewProtocol>\r\n' % (ExternalPort,Protocol)
        control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
        up = urllib2.urlparse.urlparse(ctrl_url)
        header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
        #global auth_enc
        #header += 'Authorization: Basic ' + auth_enc + '\r\n'
        header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
        header += 'CONTENT-TYPE: text/xml\r\n'
        header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
        req = header + control_doc
        host,port = up[1].split(':')
        fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        fd.connect( (host,int(port)) )
        fd.send( req )
        data = ''
        while 1:
            buf = fd.recv(1492)
            if len(buf)>0:
                data += buf
            else:
                break
        fd.close()
#        if self.is_ok(data):
#            print 'OK'
#        else:
#            print 'Fail'
        return data

    def get_port_map(self,PortMappingIndex):
        ctrl_url = self.ctrl_url
        serviceType = self.serviceType
        if ctrl_url == '':
            return ''
        actionName = 'GetGenericPortMappingEntry'
        actionParams = '<NewPortMappingIndex>%d</NewPortMappingIndex>\r\n\
<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort></NewExternalPort>\r\n\
<NewProtocol></NewProtocol>\r\n\
<NewInternalPort></NewInternalPort>\r\n\
<NewInternalClient></NewInternalClient>\r\n\
<NewEnabled>1</NewEnabled>\r\n\
<NewPortMappingDescription></NewPortMappingDescription>\r\n\
<NewLeaseDuration></NewLeaseDuration>\r\n' % (PortMappingIndex)
        control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
        up = urllib2.urlparse.urlparse(ctrl_url)
        header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
        #global auth_enc
        #header += 'Authorization: Basic ' + auth_enc + '\r\n'
        header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
        header += 'CONTENT-TYPE: text/xml\r\n'
        header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
        req = header + control_doc
        host,port = up[1].split(':')
        fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        fd.connect( (host,int(port)) )
        fd.send( req )
        data = ''
        while 1:
            buf = fd.recv(1492)
            if len(buf)>0:
                data += buf
            else:
                break
        fd.close()
        return self.analyze_map_info(data,PortMappingIndex)
    
    def analyze_map_info(self,data,i):
        #print data
        if self.is_ok(data):
            port1 = re.findall('<NewExternalPort>([0-9].*)</NewExternalPort>',data)[0]
            port2 = re.findall('<NewInternalPort>([0-9].*)</NewInternalPort>',data)[0]
            protocol = re.findall('<NewProtocol>([a-zA-Z].*)</NewProtocol>',data)[0]
            client = re.findall('<NewInternalClient>([^<>].*)</NewInternalClient>',data)[0]
            des = re.findall('<NewPortMappingDescription>([^<>].*)</NewPortMappingDescription>',data)[0]
            #print i,client,des,protocol,port1,port2
            return (i,client,des,protocol,port1,port2)
        else:
            return None
    
    def del_all_port(self):
        ctrl_url = self.ctrl_url
        serviceType = self.serviceType
        i = 0
        map_list = []
        while 1:
            try:
                ret = self.get_port_map(i)
                i += 1
                if ret != None:
                    map_list.append( ret )
                else:
                    break
            except:
                pass
        for mp in map_list:
            if mp[2]=='UPNPCoherence':
                self.del_port(int(mp[4]),mp[3])
#                print 'DEL',mp[3],mp[4]
#            else:
#                print 'Skip',mp[2]
        
    def get_ex_ip(self):
        ctrl_url = self.ctrl_url
        serviceType = self.serviceType
        if ctrl_url == '':
            return ''
        actionName = 'GetExternalIPAddress'
        actionParams = '' #'<NewRemoteHost></NewRemoteHost>\r\n\
#<NewExternalPort></NewExternalPort>\r\n\
#<NewProtocol></NewProtocol>\r\n\
#<NewInternalPort></NewInternalPort>\r\n\
#<NewInternalClient></NewInternalClient>\r\n\
#<NewEnabled>1</NewEnabled>\r\n\
#<NewPortMappingDescription></NewPortMappingDescription>\r\n\
#<NewLeaseDuration></NewLeaseDuration>\r\n'
        control_doc = '<?xml version="1.0"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\
<s:Body>\
<u:%s xmlns:u="%s">\
%s</u:%s>\
</s:Body>\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
        up = urllib2.urlparse.urlparse(ctrl_url)
        header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
        #global auth_enc
        #header += 'Authorization: Basic ' + auth_enc + '\r\n'
        header += 'CONTENT-TYPE: text/xml\r\n'
        header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n'
        header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
        header += '\r\n'
        req = header + control_doc
        host,port = up[1].split(':')
        fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        fd.connect( (host,int(port)) )
        fd.sendall( req )
        #print(req)
        data = ''
        while 1:
            buf = fd.recv(1492)
            if len(buf)>0:
                data += buf
            else:
                break
        fd.close()
        #print(data)
        self.ex_ip = re.findall('<NewExternalIPAddress>([^<>].*)</NewExternalIPAddress>',data)[0]
        return self.ex_ip

    def msearch(self):
        ssdp_req = 'M-SEARCH * HTTP/1.1\r\n\
HOST:239.255.255.250:1900\r\n\
MAN:"ssdp:discover"\r\n\
MX:3\r\n\
ST:UPnP:rootdevice\r\n\r\n'
        self.s.sendto( ssdp_req,('239.255.255.250',1900) )
        #print 'M-SEARCH send'
        
    def init_handler(self):        
        self.msearch()
        n = 0
        while 1:
            try:
                data,addr = self.s.recvfrom(1492)
                print( data )
                if len(data)>0:
                    self.location = self.get_location(data)
                    self.ctrl_url,self.serviceType = self.get_ctrlurl()
                    self.upnp=True
                    #print(self.ctrl_url,self.serviceType)
                    break
            except:
                print( 'retry msearch' )
                n+=1
                if n<5:
                    self.msearch()
                else:
                    self.upnp = False
                    break
        if self.upnp==True:
            self.get_ex_ip()
        self.get_local_ip()

def action( c,obj ):
   if c.startswith('q'):
       obj.del_all_port()
       return 1
   try:
       if c.startswith('r'):
           obj.init_handler()
           return 0
       if c.startswith('a'):
           argv = c.split()
           port_ex = int(argv[2])
           port_in = int(argv[3])
           protocol = argv[1]
           if protocol=='TCP' or protocol=='UDP':
               ret = obj.add_port(port_ex,protocol,port_in)
           else:
               print 'Unknown Protocol'
           return 0
       if c.startswith('d'):
           argv = c.split()
           port = int(argv[2])
           ret = obj.del_port(port,argv[1])
           return 0
       if c.startswith('p'):
           if c.strip()=='p':
               i = 0
               while 1:
                   try:
                       ret = obj.get_port_map(i)
                       i += 1
                       if ret != None:
                             #idx,client,des,protocol,port1,port2 = ret
                             #print idx,client,des,protocol,port1,port2
                             print ret[0],ret[1],ret[2],ret[3],ret[4],ret[5]
                       else:
                           break
                   except:
                       pass
           return 0
       return 0
       print('usage: a TCP/UDP EX_PORT IN_PORT   --add map\
                     d TCP/UDP EX_PORT    --delete map\
                     p    --show maps\
                     r    --discover device again\
                     q    --quit')
   except Exception,msg:
       print msg
       return 0

if __name__=='__main__':
    uobj = UPNPCoherence()
    uobj.init_handler()
    print 'ExternalAddress: ',uobj.ex_ip
    print 'InternalAddress: ',uobj.ip
    try:
       while 1:
           c = raw_input('>')
           ret = action( c,uobj )
           if ret==1:
            break
    except Exception,msg:
       uobj.del_all_port()
